package cn.felord.websocket.config;

import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket消息和生命周期事件的处理实现
 *
 * @author Dax
 * @since 11 :49  2018/7/6
 */
public class DefaultWebSocketHandler extends AbstractWebSocketHandler implements WebSocketHandlerProxy {
    private static final Logger log = LoggerFactory.getLogger(DefaultWebSocketHandler.class);

    /**
     * 存贮 session 对象
     */
    private static final Map<String, WebSocketSession> WEB_SOCKET_SESSIONS = new ConcurrentHashMap<>();
    private static final ObjectMapper MAPPER;

    static {
        MAPPER = new ObjectMapper();
        MAPPER.setSerializationInclusion(JsonInclude.Include.NON_NULL);
    }

    /**
     * 在WebSocket协商成功并且WebSocket连接打开并准备好使用后调用
     * 打开连接后 需要存储会话  退出后清除会话
     *
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println(session);
        String sessionId = session.getId();
        WEB_SOCKET_SESSIONS.put(sessionId, session);
        log.info("连接编号： {}  当前连接数： {}", sessionId, WEB_SOCKET_SESSIONS.size());
    }

    /**
     * 在新的WebSocket消息到达时调用
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        for (Map.Entry<String, WebSocketSession> sessionEntry : WEB_SOCKET_SESSIONS.entrySet()) {
            WebSocketSession socketSession = sessionEntry.getValue();

             System.out.println(socketSession.getAttributes());
            String sessionId = sessionEntry.getKey();
            if (socketSession.isOpen()) {
                super.handleMessage(socketSession, message);
                log.info("向 ID:{} 发送消息：{}", sessionId, message);
            } else {
                log.warn("会话: {} 已经关闭");
            }
        }
    }

    /**
     * Handle message.
     *
     * @param sessionId the session id
     * @param message   the message
     * @throws Exception the exception
     */
    @Override
    public void handleMessage(String sessionId, WebSocketMessage<?> message) throws Exception {
        WebSocketSession socketSession = WEB_SOCKET_SESSIONS.get(sessionId);

        if (socketSession != null) {
            if (socketSession.isOpen()) {
                super.handleMessage(socketSession, message);
                log.info("向 ID:{} 发送消息：{}", sessionId, message);
            } else {
                log.warn("会话: {} 已经关闭");
            }
        } else {
            log.info("webSocket session: {} 不存在", sessionId);
        }
    }

    /**
     * 文本消息处理器
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        session.sendMessage(message);
    }

    /**
     * 二进制消息处理器
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        session.sendMessage(message);
    }

    /**
     * ping-pong 消息处理器
     *
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handlePongMessage(WebSocketSession session, PongMessage message) throws Exception {
        session.sendMessage(message);
    }

    /**
     * Handle json message.
     *
     * @param <T>     the type parameter
     * @param session the session
     * @param t       the t
     * @throws Exception the exception
     */
    protected <T> void handleJsonMessage(WebSocketSession session, T t) throws Exception {
        TextMessage message = new TextMessage(MAPPER.writeValueAsBytes(t));
        handleMessage(session, message);
    }

    /**
     * Handle json message.
     *
     * @param <T>       the type parameter
     * @param sessionId the session id
     * @param t         the t
     * @throws Exception the exception
     */
    @Override
    public <T> void handleJsonMessage(String sessionId, T t) throws Exception {
        TextMessage message = new TextMessage(MAPPER.writeValueAsBytes(t));
        handleMessage(sessionId, message);
    }

    @Override
    public <T> void handleJsonMessage(T t) throws Exception {
        TextMessage message = new TextMessage(MAPPER.writeValueAsBytes(t));
        WebSocketSession socketSession = null;
        handleMessage(socketSession, message);
    }

    /**
     * 处理底层WebSocket消息传输中的错误
     *
     * @param session
     * @param exception
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {

    }

    /**
     * 在 webSocket 连接关闭后调用
     *
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        WEB_SOCKET_SESSIONS.remove(session.getId());
        log.info("会话：{} 关闭 状态码： {} 剩余连接数 {}", session.getId(), status.getCode(), WEB_SOCKET_SESSIONS.size());
    }

    /**
     * 是否支持部分消息处理  需要服务端支持  将数据量大的消息拆分 多次调用{@link WebSocketHandler#handleMessage(WebSocketSession, WebSocketMessage)} 接收
     *
     * @return
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
